import cv2
import os
from ultralytics import YOLO
from pathlib import Path
import time

def detect_and_crop_humans(
    image_path: str,
    output_dir: str = "cropped_humans",
    model_name: str = "yolov8n.pt",
    conf_threshold: float = 0.5
) -> None:
    """
    使用 YOLOv8 检测并裁剪图片中的人体
    
    参数：
    - image_path: 输入图片路径
    - output_dir: 输出目录 (默认: cropped_humans)
    - model_name: YOLO 模型名称 (默认: yolov8n)
    - conf_threshold: 置信度阈值 (默认: 0.5)
    """
    # 创建输出目录
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    # 加载 YOLO 模型
    model = YOLO(model_name)
    
    # 读取图片
    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"图片无法读取: {image_path}")
    
    # 进行目标检测
    results = model.predict(img, conf=conf_threshold, classes=[0])  # class 0 表示人体
    
    # 处理检测结果
    for i, result in enumerate(results):
        boxes = result.boxes.xyxy.cpu().numpy()
        
        # 遍历每个检测到的人体框
        for j, box in enumerate(boxes):
            x1, y1, x2, y2 = map(int, box)
            
            # 裁剪人体区域
            cropped = img[y1:y2, x1:x2]
            
            # 生成唯一文件名
            timestamp = int(time.time() * 1000)
            output_path = os.path.join(
                output_dir,
                f"{Path(image_path).stem}_human_{timestamp}_{j}.jpg"
            )
            
            # 保存图片
            cv2.imwrite(output_path, cropped)
            print(f"保存裁剪结果: {output_path}")

if __name__ == "__main__":
    # 示例用法
    detect_and_crop_humans(
        image_path="test.jpg",
        output_dir="output",
        model_name="yolov8n.pt",
        conf_threshold=0.6
    )

